# -*- coding: utf-8 -*-
import json
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling
from peft import LoraConfig, get_peft_model, PeftModel
from datasets import Dataset
from modelscope import snapshot_download


def load_dataset(dataset_path: str='./ruozhiba_qa.json', tokenizer=None):
    """
    加载数据集并预处理
    
    Args:
        dataset_path: 待训练的数据集路径
        tokenizer: 分词器，用于对文本进行 tokenize 处理
    """
    # 加载数据集
    with open(dataset_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # 转换为适合训练的格式（Dataset 格式）
    processed_data = []
    for item in data:
        # 根据模型要求调整输入输出的格式
        text = f"Human: {item['instruction']}\nAssistant: {item['output']}"
        processed_data.append({'text': text})

    # 转换为 Dataset 格式
    dataset = Dataset.from_list(processed_data)

    # 如果提供了粉刺器，进行 tokenize 处理
    if tokenizer is not None:
        dataset = dataset.map(
            lambda examples: tokenizer(
                examples['text'],
                truncation=True,
                max_length=512 # 根据需要修改
            ),
            batched=True,
            remove_columns=dataset.column_names,
        )
    return dataset


def create_load_model(base_model_path, r: int = 16, device_map: str = "auto", dtype=torch.float32):
    """
    创建 LoRA 微调模型

    Args:
        base_model_path: 基础模型路径
        r: LoRA 秩(小数据集（<5k 样本）：建议 r=8。大数据集（>50k 样本）：可尝试 r=32 甚至 r=64
        device_map: 自动设备分配（设备包括：mps-mac系统显卡、cuda-英伟达显卡、cpu）
        dtype: 精度，当显存不足时可以调为 float16。包括：float32、float16
    """
    # 加载基础模型和分词器
    tokenizer = AutoTokenizer.from_pretrained(base_model_path)
    
    # 设置填充令牌
    tokenizer.pad_token = tokenizer.eos_token
    

    # 加载基础模型
    model = AutoModelForCausalLM.from_pretrained(
        base_model_path,
        dtype=dtype, # 对 MPS 使用 float32 精度
        device_map=device_map, # 自动设备分配（设备包括：mps-mac系统显卡、cuda-英伟达显卡、cpu）
    )
    
    # 配置模型 pad token
    model.config.pad_token_id = model.config.eos_token_id

    # 冻结原始模型参数
    for param in model.parameters():
        param.requires_grad = False
    
    # LoRA 配置
    lora_config = LoraConfig(
        r=r, # LoRA 秩
        lora_alpha=32, # 缩放因子
        lora_dropout=0.1, # Dropout 比率
        bias="none", # 是否训练偏置
        task_type="CAUSAL_LM",
        target_modules=[
            "q_proj", 
            "k_proj", 
            "v_proj"
        ],
    )

    # 应用 LoRA
    pef_model = get_peft_model(model, lora_config)

    # 打印可训练参数
    pef_model.print_trainable_parameters()
    
    return pef_model, tokenizer


def train_lora_model(base_model_path, r: int = 16, device_map: str = "auto", dtype=torch.float32, dataset_path='ruozhiba_qa.json', output_dir='./lora_output'):
    """
    LoRA 微调训练
    
    Args:
        base_model_path: 基础模型路径
        r: LoRA 秩(小数据集（<5k 样本）：建议 r=8。大数据集（>50k 样本）：可尝试 r=32 甚至 r=64
        device_map: 自动设备分配（设备包括：mps-mac系统显卡、cuda-英伟达显卡、cpu）
        dtype: 精度，当显存不足时可以调为 float16。包括：float32、float16
        dataset_path: 待训练的数据集路径
        output_dir: 训练后的 LoRA 权重保存目录
    """
    # 创建 LoRA 模型
    model, tokenizer = create_load_model(base_model_path, r=r, device_map=device_map, dtype=dtype)

    # 加载数据集并 tokenize
    train_dataset = load_dataset(dataset_path, tokenizer)

    # 数据集收集器
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False, # 因为是因果语言模型，所以不使用 MLM
    )

    # 训练参数
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=3,
        per_device_train_batch_size=1, # 在 MPS 上可能需要更小的 batch size
        gradient_accumulation_steps=4, # 模拟更大的 batch size
        learning_rate=1e-4,
        logging_steps=10,
        remove_unused_columns=False,
        push_to_hub=False,
        logging_dir='./logs',
)

    # 训练器
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        data_collator=data_collator,
    )

    # 开始训练
    trainer.train()

    # 保存 LoRA 权重
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)


def inference_lora_model(base_model_path, lora_path='./lora_output', dtype=torch.float32):
    """
    使用 LoRA 微调后的模型进行推理

    Args:
        base_model_path: 基础模型路径
        lora_path: LoRA 训练后的权重路径
        dtype: 精度，当显存不足时可以调为 float16。包括：float32、float16
    """
    # 加载原始模型
    model = AutoModelForCausalLM.from_pretrained(
        base_model_path,
        dtype=dtype, # 使用 float32 精度
    )
    tokenizer = AutoTokenizer.from_pretrained(base_model_path)
    
    # 加载 LoRA 权重
    model = PeftModel.from_pretrained(model, lora_path)

    # 合并 LoRA 权重到原始模型
    model = model.merge_and_unload()
    
    # 设置填充令牌
    tokenizer.pad_token = tokenizer.eos_token
    model.config.pad_token_id = model.config.eos_token_id

    # 测试推理
    test_prompts = [
        "只剩一个心脏了还能活吗？",
        "如何知道学习一门编程语言？",
        "你是谁？能帮我解决什么问题？",
    ]

    for prompt in test_prompts:
        # 构造完整的输入
        full_prompt = f"Human: {prompt}\nAssistant:"

        inputs = tokenizer(full_prompt, return_tensors="pt")

        outputs = model.generate(
            **inputs,
            max_length=100,
            num_return_sequences=1,
            temperature=0.7,
        )

        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        print(f"提示：{prompt}")
        print(f"回复：{generated_text}\n")


def merge_and_save_lora_model(base_model_path, lora_path='./lora_output', dtype=torch.float32, save_path='./model_output'):
    """
    合并 LoRA 权重到基础模型并保存为完整模型
    
    Args:
        base_model_path: 基础模型路径
        lora_path: LoRA 训练后的权重路径
        dtype: 精度，当显存不足时可以调为 float16。包括：float32、float16
        save_path: 合并后模型的保存路径
    """
    # 加载原始模型
    model = AutoModelForCausalLM.from_pretrained(
        base_model_path,
        dtype=dtype,
    )
    tokenizer = AutoTokenizer.from_pretrained(base_model_path)
    
    # 加载 LoRA 权重
    model = PeftModel.from_pretrained(model, lora_path)

    # 合并 LoRA 权重到原始模型
    model = model.merge_and_unload()
    
    # 设置填充令牌
    tokenizer.pad_token = tokenizer.eos_token
    model.config.pad_token_id = model.config.eos_token_id

    # === 保存合并后的模型 ===
    print(f"开始保存合并后的模型到 {save_path}")
    model.save_pretrained(save_path)
    tokenizer.save_pretrained(save_path)
    
    # 可选：保存为 safetensors 格式（更安全）
    model.save_pretrained(
        save_path, 
        safe_serialization=True  # 这会生成 .safetensors 文件
    )
    tokenizer.save_pretrained(save_path)
    print("模型保存完成！")


def download(model_id: str = "LLM-Research/Llama-3.2-1B"):
    """下载基础模型"""
    # 从ModelScope下载模型--第一次比较耗时
    base_model_path = snapshot_download(
        model_id=model_id,
        cache_dir="./model_cache",
    )


if __name__ == "__main__":
    model_id="LLM-Research/Llama-3.2-1B" # 模型 id，可以从 https://modelscope.cn/ 查找对应 model_id
    base_model_path = "./model_cache/LLM-Research/Llama-3___2-1B" # 基础模型下载后保存的本地路径

    dataset_path = "./ruozhiba_qa.json" # 待训练的数据集路径
    r=16 # LoRA 秩(小数据集（<5k 样本）：建议 r=8。大数据集（>50k 样本）：可尝试 r=32 甚至 r=64
    device_map="mps" # 自动设备分配（设备包括：auto-自动检测，不要使用auto，可能回报错、mps-mac系统显卡、cuda-英伟达显卡、cpu）
    dtype=torch.float32 # 精度，当显存不足时可以调为 float16。包括：float32、float16

    # 1、下载基础模型
    # download(model_id)

    # 2、训练 LoRA 模型
    train_lora_model(base_model_path, r=r, device_map=device_map, dtype=dtype, dataset_path=dataset_path)

    # 3、使用 LoRA 微调后的模型进行推理（测试）
    # inference_lora_model(base_model_path, dtype=dtype)

    # 4、合并 LoRA 权重到基础模型并保存为完整模型
    # merge_and_save_lora_model(base_model_path, dtype=dtype)
